import numpy as np
from scipy.fftpack import dst, idst

# 定义泊松方程的解析解
def exact_solution(X, Y):
    return  np.sin(np.pi * X) * np.sin(2 * np.pi * Y)

# 定义泊松方程的源项
def source_term(X, Y):
    return -5 * np.pi**2 * np.sin(np.pi * X) * np.sin(2 * np.pi * Y)

# 定义解二维泊松方程的函数
def poisson_2d_periodic_x_dirichlet_y(Nx, Ny):
    # 创建网格
    X = np.linspace(0, 1, Nx, endpoint=False)
    Y = np.linspace(0, 1, Ny, endpoint=True)  # y方向的最后一个点是1，因为我们应用Dirichlet边界条件
    X, Y = np.meshgrid(X, Y)
    
    # 源项f
    f = source_term(X, Y)
    
    # y方向的DST-II变换
    F_dst = dst(f, type=2, axis=0)
    
    # 由于周期性边界条件，x方向使用FFT
    F_fft = np.fft.fft2(f)
    
    # 构造k空间中的泊松方程
    kx = 2 * np.pi * np.fft.fftfreq(Nx, d=1/Nx)
    ky = 2 * np.pi * np.fft.fftfreq(Ny, d=1/Ny)
    K2 = np.outer(kx, np.ones(Ny))**2 + np.outer(np.ones(Nx), ky)**2
    U_k = -F_fft / K2
    
    # 由于DST已经考虑了边界条件，我们需要将U_k的边界设置为0
    U_k[:, 0] = U_k[:, -1] = 0
    
    # y方向的逆DST-II变换
    psi_dst = idst(U_k, type=2, axis=0)
    
    # x方向的逆FFT
    psi_fft = np.fft.ifft2(psi_dst)
    
    # 应用y方向的Dirichlet边界条件
    psi_fft[:, 0] = psi_fft[:, -1] = 0
    
    return psi_fft

# 选择网格点数目
Nx = 64
Ny = 64

# 计算数值解
u_num = poisson_2d_periodic_x_dirichlet_y(Nx, Ny)

# 计算解析解
u_exact = exact_solution(np.linspace(0, 1, Nx, endpoint=False), np.linspace(0, 1, Ny, endpoint=True))

# 计算误差
error = np.abs(u_num - u_exact)

# 显示结果
print("最大误差:", np.max(error))
print("均方误差:", np.mean(error**2))